from torch.utils.tensorboard import SummaryWriter
import os
from tensorboard.backend.event_processing import event_accumulator
"""
记录学习数据脚本
"""

def record_train_value(loss_value, accuracy, lr, epoch):
    """

    :param loss_value: 损失值
    :param accuracy: 精确度
    :param lr: 学习率
    :param epoch: 当前迭代次数
    :return:
    """

def init_tensorboard(name):
    if not os.path.exists(name):
        os.makedirs(name)
    return SummaryWriter(name)

def visual_model_tensorboard_script(writer, model, images):
    writer.add_graph(model, images)

def visual_scalars_tensorboard_script(writer, tag, value, epoch):
    writer.add_scalar(tag, value, epoch)

def close_tensorboard(writer):
    writer.close()

def read_tensorboard():
    """
    读取tensorboard数据
    :return:
    """
    # 加载日志数据
    ea = event_accumulator.EventAccumulator(os.path.join(os.getcwd(), "../runs/train"))
    ea.Reload()
    print(ea.scalars.Keys())

    val_acc = ea.scalars.Items('val_acc')
    print(len(val_acc))
    print([(i.step, i.value) for i in val_acc])

    import matplotlib.pyplot as plt
    fig = plt.figure(figsize=(6, 4))
    ax1 = fig.add_subplot(111)
    val_acc = ea.scalars.Items('val_acc')
    ax1.plot([i.step for i in val_acc], [i.value for i in val_acc], label='val_acc')
    ax1.set_xlim(0)
    acc = ea.scalars.Items('acc')
    ax1.plot([i.step for i in acc], [i.value for i in acc], label='acc')
    ax1.set_xlabel("step")
    ax1.set_ylabel("")

    plt.legend(loc='lower right')
    plt.show()